
import torch

from torch import nn

# import pdb
import numpy as np
import random
# from . import networks
from .base_model import BaseModel
from .tianyu import network

from PIL import Image
from data.base_dataset import  get_transform
from torchvision import transforms
import cv2
from .tianyu.testapi import TestApi

def deconv_orth_dist(kernel, stride = 2, padding = 1,device=None):
    [o_c, i_c, w, h] = kernel.shape
    output = torch.conv2d(kernel, kernel, stride=stride, padding=padding)
    target = torch.zeros((o_c, o_c, output.shape[-2], output.shape[-1])).to(device)
    ct = int(np.floor(output.shape[-1]/2))
    if device is None:
        target[:,:,ct,ct] = torch.eye(o_c).cuda()
    else:
        target[:,:,ct,ct] = torch.eye(o_c).to(device)
    return torch.norm( output - target )
    
def orth_dist(mat, stride=None,device=None):
    mat = mat.reshape( (mat.shape[0], -1) )
    if mat.shape[0] < mat.shape[1]:
        mat = mat.permute(1,0)
    if device is None:
        return torch.norm( torch.t(mat)@mat - torch.eye(mat.shape[1]).cuda())
    else:
        return torch.norm( torch.t(mat)@mat - torch.eye(mat.shape[1]).to(device))

def one_hot(batchsize,numclass,tensor):
    return torch.zeros(batchsize,numclass).to(tensor.device).scatter_(1,tensor.view(1,tensor.shape[0]).long().t(),1)

class polarTianYuModel(BaseModel):
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

        Returns:
            the modified parser.

        
        """
        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
        parser.add_argument('--num_classes', type=int, default=3, help='the architecture situation of model')
        parser.add_argument('--arch_name', type=str, default='mobilenet', help='the arch of network')
        parser.add_argument('--arch_type', type=str, default='', help='the arch of network')
        parser.add_argument('--pretrained', type=bool, default=False, help='the arch of network')
        parser.add_argument('--input_modality', type=str, default='SD', help='the modality of input data')
        parser.add_argument('--lambda_orth', type=float, default=0.1, help='the modality of input data')
        # parser.add_argument('--init_type', type=str, default='norm', help='the init type of network')
        # parser.add_argument('--dataIn', type=str, default='s0', help='the data type of input')
        # parser.add_argument('--label', type=str, default='dolp', help='the data type of label')

            # parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')
            # parser.add_argument('--lambda_identity', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')

        return parser
    def display_opt_init(opt):
        # parser.add_argument('--inputs', type=str, default='img', help='the moudal of input data ')
        opt.__setattr__('num_classes',2)
        opt.__setattr__('arch_name','mobilenet')
        opt.__setattr__('arch_type','')
        opt.__setattr__('pretrained',False)
        opt.__setattr__('input_modality','SD')
        return opt

    # def display_init(self,_dict):
    #     # transform = get_transform(self.opt, grayscale=(self.opt.input_nc == 1))
        
    #     # if isinstance(_dict['img'],np.ndarray):
    #     #     _dict[self.opt.input_modality] = None
    #     # else:
    #     #     _dict[self.opt.input_modality] = _dict['img']
    #     return _dict

    def __init__(self, opt):#basic_conv=networks.Conv2d_cd, theta=0.7 ):   
        BaseModel.__init__(self, opt)
        
        self.testapi = TestApi(opt)
        self.model_names = ['polarLD']
        self.netpolarLD= network.define_network(opt)
        # getattr(network,opt.arch_name)(opt)
        
        self.visuale_param()
        
        if self.isTrain:
            self.criterionCM = nn.L1Loss()
            self.criterion = nn.L1Loss()
            self.optimizer = torch.optim.Adam(self.netpolarLD.parameters(), lr=opt.lr, weight_decay=1e-3)
            self.optimizers.append(self.optimizer)
        # self.set_requires_grad([self.netpolarLD], False)
            
    def visuale_param(self):
        self.loss_names = ["Accurate","F_score","APCER","NPCER","ACER"]
        self.testapi.init(self.opt.controller,["Accurate","F_score","APCER","NPCER","ACER"])
        if self.isTrain:
            self.loss_names += ["polarLD","orth"]
        elif self.opt.controller != 'test':
                self.last_val_best["Accurate"] = 0
        # self.last_val = {}
        # for key in self.loss_names:
        #     self.last_val[key] = 0
        
            # if opt.lambda_identity > 0.0:  # only works when input and output images have the same number of channels
            #     assert(opt.input_nc == opt.output_nc)
            # self.fake_A_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            # self.fake_B_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            # define loss functions
            
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            
            # self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
            # self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
            
            # self.optimizers.append(self.optimizer_D)
        pass
    def set_input(self, data):
        if "isTrain" in data.keys():
            if self.isTrain != data["isTrain"]:
                self.isTrain = data["isTrain"]
                self.visuale_param()
                
        if self.opt.input_modality in data.keys() and 'G'+self.opt.input_modality in data.keys():
            indata = torch.cat((data[self.opt.input_modality],data['G'+self.opt.input_modality],))

            label = torch.cat((torch.zeros(data[self.opt.input_modality].shape[0])+1 ,torch.zeros(data['G'+self.opt.input_modality].shape[0]),),0)
            
            ind = np.arange(0,label.shape[0])
            random.shuffle(ind)
            
            self.data = indata[ind].to(self.device)
            self.label = label[ind]
            self.label_1hot = one_hot(self.label.shape[0],self.opt.num_classes,self.label).to(self.device)
            
        elif self.opt.input_modality in data.keys() and 'label' in data.keys():
            self.data = data[self.opt.input_modality].to(self.device)
            self.label = data['label']
            self.label_1hot = one_hot(self.label.shape[0],self.opt.num_classes,self.label).to(self.device)
        
        elif self.opt.input_modality in data.keys():
            self.data = data[self.opt.input_modality]
            if self.data is not None:
                self.data = self.data.to(self.device)
        else:
            self.data = None
        self.image_paths = None
        
        
        
        if not self.isTrain:
            if 'cat_name' in data.keys():
                    self.cat_name = data['cat_name']
            else:
                self.cat_name = {'fake':0,'real':1}
            
            if 'category' in data.keys():
                self.cat= data['category']
            elif  'G'+self.opt.input_modality in data.keys() :
                self.cat = label
                
        pass
    def forward(self):	    	# x [3, 256, 256]
        if self.data is not None:
            # B,C,H,W = self.data.shape
            # self.out = torch.zeros((B,2)).to(self.device)
            # for i in range(B):
            #     self.out[i,:] = self.netpolarLD(self.data[i:i+1,:,:,:])['prediction']
            self.out = self.netpolarLD(self.data)
                
        # self.out,self.fc_vis = self.netpolarLD(self.data['S0'])
        # self.out_s1,self.fc_s1 = self.netpolarLD(self.data['S1'])
        # self.out_s2,self.fc_s2 = self.netpolarLD(self.data['S2'])
        pass
    def backward(self):
        # self.loss_cross_modal = self.opt.lambda_cm * self.criterionCM(self.fc['fc'][self.s0len:self.s0len+self.s1len],self.fc['fc_CM'][0:self.s0len])
        self.loss_polarLD = self.criterion(self.out['prediction'],self.label_1hot)
        if self.opt.init_type == 'orthogonal':
            if self.opt.arch_name == 'cdcn':
                self.loss_orth = orth_dist(self.netpolarLD.module.fc2.weight,device=self.opt.gpu_ids[0])
                # self.loss_orth += orth_dist(self.netpolarLD.module.downsample32x32[0].weight,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.Block1[0].conv.weight, stride=1,device=self.opt.gpu_ids[0]) + deconv_orth_dist(self.netpolarLD.module.sa1.conv1.weight, stride=1,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.Block2[0].conv.weight, stride=1,device=self.opt.gpu_ids[0]) + deconv_orth_dist(self.netpolarLD.module.sa2.conv1.weight, stride=1,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.Block3[0].conv.weight, stride=1,device=self.opt.gpu_ids[0]) + deconv_orth_dist(self.netpolarLD.module.sa3.conv1.weight, stride=1,device=self.opt.gpu_ids[0])

            elif self.opt.arch_name == 'resnet':
                self.loss_orth = orth_dist(self.netpolarLD.module.classifier.classifier.weight,device=self.opt.gpu_ids[0])
                self.loss_orth += orth_dist(self.netpolarLD.module.features[5][0].downsample[0].weight,device=self.opt.gpu_ids[0]) + orth_dist(self.netpolarLD.module.features[6][0].downsample[0].weight,device=self.opt.gpu_ids[0]) + orth_dist(self.netpolarLD.module.features[7][0].downsample[0].weight,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.features[4][0].conv1.weight, stride=1,device=self.opt.gpu_ids[0]) + deconv_orth_dist(self.netpolarLD.module.features[4][1].conv1.weight, stride=1,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.features[5][0].conv1.weight, stride=2,device=self.opt.gpu_ids[0]) + deconv_orth_dist(self.netpolarLD.module.features[5][1].conv1.weight, stride=1,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.features[6][0].conv1.weight, stride=2,device=self.opt.gpu_ids[0]) + deconv_orth_dist(self.netpolarLD.module.features[6][1].conv1.weight, stride=1,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.features[7][0].conv1.weight, stride=2,device=self.opt.gpu_ids[0]) + deconv_orth_dist(self.netpolarLD.module.features[7][1].conv1.weight, stride=1,device=self.opt.gpu_ids[0])
            elif self.opt.arch_name == 'vgg':
                self.loss_orth = orth_dist(self.netpolarLD.module.classifier.classifier.weight,device=self.opt.gpu_ids[0])
                # self.loss_orth += orth_dist(self.netpolarLD.module.features[5][0].downsample[0].weight,device=self.opt.gpu_ids[0]) + orth_dist(self.netpolarLD.module.features[6][0].downsample[0].weight,device=self.opt.gpu_ids[0]) + orth_dist(self.netpolarLD.module.features[7][0].downsample[0].weight,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.features[7].weight, stride=1,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.features[14].weight, stride=1,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.features[24].weight, stride=1,device=self.opt.gpu_ids[0])
                self.loss_orth += deconv_orth_dist(self.netpolarLD.module.features[34].weight, stride=1,device=self.opt.gpu_ids[0])
            self.loss_polarLD = self.loss_polarLD + self.opt.lambda_orth * self.loss_orth
            # diff = orth_dist(self.netpolarLD.module.features[5][0].downsample[0].weight) + orth_dist(self.netpolarLD.module.features[6][0].downsample[0].weight) + orth_dist(self.netpolarLD.module.features[7][0].downsample[0].weight)
            # diff += deconv_orth_dist(self.netpolarLD.module.features[4][0].conv1.weight, stride=1) + deconv_orth_dist(self.netpolarLD.module.features[4][1].conv1.weight, stride=1)
            # diff += deconv_orth_dist(self.netpolarLD.module.features[5][0].conv1.weight, stride=2) + deconv_orth_dist(self.netpolarLD.module.features[5][1].conv1.weight, stride=1)
            # diff += deconv_orth_dist(self.netpolarLD.module.features[6][0].conv1.weight, stride=2) + deconv_orth_dist(self.netpolarLD.module.features[6][1].conv1.weight, stride=1)
            # diff += deconv_orth_dist(self.netpolarLD.module.features[7][0].conv1.weight, stride=2) + deconv_orth_dist(self.netpolarLD.module.features[7][1].conv1.weight, stride=1)
        #####
        # self.loss = self.loss_polarLD
        self.loss_polarLD.backward()
        # self.cal_score()
    def optimize_parameters(self):
        self.forward()
        self.optimizer.zero_grad()
        self.backward()
        self.optimizer.step()
        pass
    
    # def cal_score(self):
    def iter_end(self):
        # self.loss_pred_acc = (torch.where(ans == self.label)[0].shape[0])  / self.label.shape[0]
        # self.count_mean_pred_acc.update(self.loss_pred_acc,self.label.shape[0])
        # self.loss_mean_pred_acc = self.count_mean_pred_acc.avg
        # if self.opt.controller == 'test' or not self.isTrain:
        # label,out,controller,cat,cat_name
        records = {}
        records['label'] = self.label
        records['out'] = self.out
        records['controller'] = self.opt.controller
        if self.opt.controller == 'test' or not self.isTrain:
            records['cat'] = self.cat
            records['cat_name'] = self.cat_name
        param = self.testapi.record(records)

        self.loss_Accurate = param['Accurate']['all']
        self.loss_F_score = param['F_score']['all']
        self.loss_APCER = param['APCER']['all']
        self.loss_NPCER = param['NPCER']['all']
        self.loss_ACER = param['ACER']['all']
        self.loss_MACER = 1 - self.loss_ACER

        pass

    def epoch_end(self):
        # self.count_mean_pred_acc.reset()
        # if not self.isTrain:
        self.testapi.statistic(self.opt.controller)
            # param = get_res(self.TP,self.TN,self.FP,self.FN,2)
            
            # if self.opt.controller == 'test':
            #     csvdir = os.path.join(self.opt.results_dir,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch)
            #     csvfile = '{}_{}_{}_{}'.format(self.opt.arch_name,self.opt.arch_type,self.opt.input_modality,'v'+self.opt.version)
            #     sheet_name = '{}{}'.format(self.opt.arch_name+self.opt.arch_type,self.opt.input_modality.upper())
            #     if not '物理' in self.TP.keys():
            #         save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,list(self.cat_name.keys())+['all','物理'])
            #     elif self.cat_name is not None:
            #         save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,list(self.cat_name.keys())+['all'])
            #     else:
            #         save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,['all'])

            # # roc_plot(self.groud_truth,self.prediction,'test.png','/home/yalin.huang/test/results/'+self.opt.name)
            #     mergetotxt(self.groud_truth,self.prediction,os.path.join(os.path.abspath('.'),'results',self.opt.name),self.opt.arch_name,self.opt.version)
            #     self.TP['物理'],self.TN['物理'],self.FP['物理'],self.FN['物理'], = 0,0,0,0
            # self.TP,self.TN,self.FP,self.FN = {},{},{},{}
            # self.TP['all'],self.TN['all'],self.FP['all'],self.FN['all']  = 0,0,0,0
        pass
    
    def val_save(self, param):
        if getattr(self,'loss_' + param) == 0 and self.last_val_best[param] == 0:
            if self.last_val_best["Accurate"] < self.loss_Accurate:
                self.last_val_best["Accurate"] = self.loss_Accurate
                return True
        else:
            if self.last_val_best[param] < getattr(self,'loss_' + param):
                self.last_val_best[param] = getattr(self,'loss_' + param)
                return True
        return False
        pass
    
    def display_proc(self,_dict):
        tran = transforms.ToPILImage()
        if len(_dict['boxes'])>0:
            preds = torch.max(self.out,axis=1).indices.cpu()
            for i,pred in enumerate(preds):
                if pred.item():
                    cv2.rectangle(_dict['box_img'], (_dict['boxes'][i][0], _dict['boxes'][i][1]), (_dict['boxes'][i][2], _dict['boxes'][i][3]), (0, 0, 255), 2)
                    text = 'real'
                    cx = _dict['boxes'][i][0]
                    cy = _dict['boxes'][i][1] + 12
                    cv2.putText(_dict['box_img'], text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
                else:
                    cv2.rectangle(_dict['box_img'], (_dict['boxes'][i][0], _dict['boxes'][i][1]), (_dict['boxes'][i][2], _dict['boxes'][i][3]), (255, 0, 0), 2)
                    text = 'fake'
                    cx = _dict['boxes'][i][0]
                    cy = _dict['boxes'][i][1] + 12
                    cv2.putText(_dict['box_img'], text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))

            _dict['img'] = _dict['box_img']
        else:
                _dict['img'] = _dict['ori_img']
                
        return _dict